from torch.utils.data import DataLoader
from dataset import TrafficDataset
from model import TransAm, LSTMModel, lstm_init_weights, Transformer
import argparse
import torch
import torch.nn as nn
from tqdm import tqdm
import math
from utils import PlotUtils
from constant import Constant


def get_model(args):
    feature_choose = Constant.FEATURES_CHOOSE
    n_inputs = len(Constant.FEATURES_MAP[feature_choose]) + 1

    model_name = Constant.MODELS[Constant.MODEL_CHOOSE]
    print(f'you are using {model_name}...')
    if model_name == 'transformer_encoder':
        return TransAm(input_feature_size=n_inputs).to(args.device)
    elif model_name == 'LSTM':
        return LSTMModel(n_input=n_inputs).to(args.device).apply(lstm_init_weights)
    elif model_name == 'transformer':
        args.have_decoder = True
        return Transformer(n_encoder_inputs=n_inputs, n_decoder_inputs=n_inputs).to(args.device)
    else:
        raise ValueError


# TODO 多步预测
"""使用当前模型预测输出作为下一次预测的输入，从而进行多步的预测，区别于OUT_WINDOW>1的情况"""
